import os
import json
import re
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
from src.llamafactory.chat.chat_model import ChatModel
from tqdm import tqdm
import argparse
import json

output = '/home/data/saves/sft'
parser = argparse.ArgumentParser(description="Your script description")
parser.add_argument("--model_name_or_path", default="Qwen2-7B-Instruct", type=str,
                    help="Model name or path")
parser.add_argument("--adapter_name_or_path", default=output, type=str, help="Checkpoint directory")
parser.add_argument("--template", default="qwen", type=str, help="Template type")
parser.add_argument("--finetuning_type", default="lora", type=str, help="Fine-tuning type")

args = parser.parse_args()
args_dict = vars(args)
chat_model = ChatModel(args_dict)
history = []

with open('data/test.jsonl','r') as f:
    data = f.readlines()

f= open(f"{output}/submission.json", "w", encoding="utf-8")

num=0
for d in tqdm(data):
    d = json.loads(d)
    apis = d['apis']
    user_message = d['user_messages']

    prompt = f"""你是一名专业的智能体（agent）专家。在智能体执行任务时，需要依据用户多轮对话内容进行总结，明确所需完成的问题，并确定使用哪些工具API以及对应的参数。
    ##要求：
    1.用户的多轮对话中只有一个工具api被调用，需要给出工具名
    2.根据多轮对话内容，判断出在第几轮可以满足工具调用的参数需求
    
    输出正确的结果：
    apis:
    {apis}
    用户多轮对话内容:
    {user_message} 
    """
   
    messages = [
        {"role": "user", "content":   prompt}]

    response = chat_model.chat(messages, history=[],  do_sample=False,num_beams=3)
    
    res = response[0].response_text.strip()
    # 把res从字符串转换成list
    try:
        # res = ast.literal_eval(res)
        res = eval(res.replace('true','True').replace('false','False').replace('null','None'))
        assert isinstance(res, list)
    except:
        print(res)
        res = []
    d['targets']=res
    f.write(json.dumps(d, ensure_ascii=False) + "\n")
    f.flush()